#include "Optimizer.h"
#include "CoreLogic.h"
#include <iostream>
#include <random>
#include <algorithm>
#include <vector>
#include <omp.h>

OptimizerNSGA::OptimizerNSGA(const std::vector<int>& initial_schedule, const Graph& graph,
    const CodeMap& code_map, const Caps& caps, const Params& initial_params,
    const Options& options)
    : initial_schedule_(initial_schedule), graph_(graph), code_map_(code_map),
    caps_(caps), initial_params_(initial_params), options_(options) {
}

void OptimizerNSGA::precomputeAncestors() {
    std::cout << "  Optimizer > Precomputing ancestor matrix..." << std::endl;
    int num_nodes = graph_.num_nodes;
    is_ancestor_.assign(num_nodes + 1, std::vector<bool>(num_nodes + 1, false));
    for (int i = 1; i <= num_nodes; ++i) {
        std::vector<int> q = graph_.predecessors_list[i];
        std::vector<bool> visited(num_nodes + 1, false);
        int head = 0;
        while (head < q.size()) {
            int curr = q[head++];
            if (!visited[curr]) {
                visited[curr] = true;
                is_ancestor_[curr][i] = true;
                for (int pred : graph_.predecessors_list[curr]) {
                    q.push_back(pred);
                }
            }
        }
    }
    std::cout << "  Optimizer > Precomputation complete." << std::endl;
}

std::vector<Individual> OptimizerNSGA::run() {
    std::cout << "--- Starting multi-objective optimization (optimizer_v_nsga) ---" << std::endl;

    precomputeAncestors();

    std::cout << "  Optimizer > Evaluating initial solution..." << std::endl;
    auto [initial_time, initial_spill, initial_best_alloc_strategy, initial_final_params] =
        evaluateCandidateThreeWays(initial_schedule_, initial_params_, {});

    if (std::isinf(initial_time)) {
        throw std::runtime_error("Optimization failed: Initial solution timed out.");
    }

    initial_scores_ = { initial_time, std::max(1.0, initial_spill) };
    std::cout << "  Optimizer > Benchmark -> Time: " << initial_time << ", Spill: " << initial_spill
        << " (using " << initial_best_alloc_strategy << " strategy)" << std::endl;

    std::cout << "  Optimizer > Creating and evaluating initial population..." << std::endl;
    std::vector<Individual> population;
    population.emplace_back(initial_schedule_, initial_final_params, initial_time, initial_spill);

    std::vector<Individual> initial_pop_additions(options_.population_size - 1);

#pragma omp parallel for
    for (int i = 0; i < options_.population_size - 1; ++i) {
        std::mt19937 rng(std::random_device{}() + omp_get_thread_num());
        std::uniform_real_distribution<> dist(0.0, 1.0);

        auto temp_schedule = initial_schedule_;
        auto temp_params = initial_params_;
        if (dist(rng) < 0.5) {
            temp_schedule = mutateSchedule(temp_schedule);
        }
        else {
            temp_params = mutateParams(temp_params);
        }
        auto [t, s, as, p] = evaluateCandidateThreeWays(temp_schedule, temp_params, initial_scores_);
        initial_pop_additions[i] = Individual(temp_schedule, p, t, s);
    }
    population.insert(population.end(), initial_pop_additions.begin(), initial_pop_additions.end());

    nonDominatedSort(population);

    for (int gen = 1; gen <= options_.generations; ++gen) {
        std::cout << "  Optimizer > --- Generation " << gen << "/" << options_.generations << " ---" << std::endl;

        std::vector<int> front1_indices;
        for (int i = 0; i < population.size(); ++i) if (population[i].rank == 1) front1_indices.push_back(i);

        if (front1_indices.empty()) {
            std::cerr << "Warning: No solutions in Front 1. Using whole population as parents." << std::endl;
            for (int i = 0; i < population.size(); ++i) front1_indices.push_back(i);
        }

        std::vector<Individual> offspring(options_.batch_size);

#pragma omp parallel for
        for (int k = 0; k < options_.batch_size; ++k) {
            std::mt19937 rng(std::random_device{}() + omp_get_thread_num());
            std::uniform_real_distribution<> dist(0.0, 1.0);
            std::uniform_int_distribution<> parent_dist(0, front1_indices.size() - 1);

            const auto& parent = population[front1_indices[parent_dist(rng)]];

            std::vector<int> child_schedule;
            Params child_params;

            if (dist(rng) < options_.schedule_move_prob) {
                child_schedule = mutateSchedule(parent.schedule);
                child_params = parent.params;
            }
            else {
                child_schedule = parent.schedule;
                child_params = mutateParams(parent.params);
            }
            auto [t, s, as, p] = evaluateCandidateThreeWays(child_schedule, child_params, initial_scores_);
            offspring[k] = Individual(child_schedule, p, t, s);
        }

        std::vector<Individual> combined_population = population;
        combined_population.insert(combined_population.end(), offspring.begin(), offspring.end());

        nonDominatedSort(combined_population);

        std::vector<Individual> new_population;
        int front_num = 1;
        while (true) {
            std::vector<Individual> front;
            for (const auto& ind : combined_population) if (ind.rank == front_num) front.push_back(ind);

            if (new_population.size() + front.size() <= options_.population_size) {
                new_population.insert(new_population.end(), front.begin(), front.end());
            }
            else {
                std::sort(front.begin(), front.end(), [](const Individual& a, const Individual& b) {
                    return a.crowding_distance > b.crowding_distance;
                    });
                size_t remaining_space = options_.population_size - new_population.size();
                new_population.insert(new_population.end(), front.begin(), front.begin() + remaining_space);
                break;
            }
            front_num++;
            bool front_exists = false;
            for (const auto& ind : combined_population) if (ind.rank == front_num) front_exists = true;
            if (!front_exists) break;
        }
        population = new_population;

        int current_front1_count = 0;
        for (const auto& p : population) if (p.rank == 1) current_front1_count++;
        std::cout << "  Optimizer > Current pareto front has " << current_front1_count << " solutions." << std::endl;
    }

    std::vector<Individual> pareto_front;
    for (const auto& p : population) if (p.rank == 1) pareto_front.push_back(p);

    std::cout << "--- Multi-objective optimization complete ---" << std::endl;
    return pareto_front;
}


std::tuple<double, double, std::string, Params> OptimizerNSGA::evaluateCandidateThreeWays(
    const std::vector<int>& schedule, const Params& spill_params,
    const std::optional<std::pair<double, double>>& initial_scores) {

    const std::vector<std::string> strategies = { "BestFit", "LowAddress", "LifetimeAware" };
    double best_time = std::numeric_limits<double>::infinity();
    double best_spill = std::numeric_limits<double>::infinity();
    std::string best_alloc_strategy = "None";

    for (const auto& strat : strategies) {
        Params current_params = spill_params;
        current_params.alloc_strategy = strat;

        auto [time, spill] = CoreLogic::evaluateSchedule(schedule, graph_, code_map_, caps_, current_params, initial_scores);

        if ((time < best_time && spill <= best_spill) || (time <= best_time && spill < best_spill)) {
            best_time = time;
            best_spill = spill;
            best_alloc_strategy = strat;
        }
        else if (time == best_time && spill == best_spill) {
            best_alloc_strategy = strat;
        }
    }

    Params final_params = spill_params;
    final_params.alloc_strategy = best_alloc_strategy;
    return { best_time, best_spill, best_alloc_strategy, final_params };
}


std::vector<int> OptimizerNSGA::mutateSchedule(std::vector<int> schedule) const {
    if (schedule.empty()) return schedule;
    std::mt19937 rng(std::random_device{}());
    std::uniform_int_distribution<> dist(0, schedule.size() - 1);

    int node_idx_to_move = dist(rng);
    int node_to_move = schedule[node_idx_to_move];

    int latest_ancestor_pos = -1;
    for (int i = 0; i < schedule.size(); ++i) {
        if (is_ancestor_[schedule[i]][node_to_move]) {
            latest_ancestor_pos = i;
        }
    }

    int earliest_descendant_pos = schedule.size();
    for (int i = schedule.size() - 1; i >= 0; --i) {
        if (is_ancestor_[node_to_move][schedule[i]]) {
            earliest_descendant_pos = i;
        }
    }

    std::vector<int> valid_positions;
    for (int i = latest_ancestor_pos + 1; i < earliest_descendant_pos; ++i) {
        if (i != node_idx_to_move) {
            valid_positions.push_back(i);
        }
    }

    if (!valid_positions.empty()) {
        std::uniform_int_distribution<> pos_dist(0, valid_positions.size() - 1);
        int new_pos_in_valid = pos_dist(rng);
        int new_schedule_pos = valid_positions[new_pos_in_valid];

        schedule.erase(schedule.begin() + node_idx_to_move);
        if (new_schedule_pos > node_idx_to_move) new_schedule_pos--;
        schedule.insert(schedule.begin() + new_schedule_pos, node_to_move);
    }

    return schedule;
}

Params OptimizerNSGA::mutateParams(Params params) const {
    std::mt19937 rng(std::random_device{}());
    std::uniform_real_distribution<> dist(-0.5, 0.5);

    double thr_range = options_.threshold_range.second - options_.threshold_range.first;
    double mutation = dist(rng) * thr_range * 0.1;
    params.min_farthest_use_threshold += mutation;
    params.min_farthest_use_threshold = std::max(options_.threshold_range.first,
        std::min(options_.threshold_range.second, params.min_farthest_use_threshold));

    std::uniform_real_distribution<> weight_dist(-0.1, 0.1);
    double w1 = params.w_use + weight_dist(rng);
    double w2 = params.w_vol + weight_dist(rng);
    double w3 = params.w_vic + weight_dist(rng);

    double sum = std::fabs(w1) + std::fabs(w2) + std::fabs(w3);
    if (sum > 0) {
        params.w_use = std::fabs(w1) / sum;
        params.w_vol = std::fabs(w2) / sum;
        params.w_vic = std::fabs(w3) / sum;
    }

    return params;
}

void OptimizerNSGA::nonDominatedSort(std::vector<Individual>& population) const {
    std::vector<std::vector<int>> fronts(1);

    for (int i = 0; i < population.size(); ++i) {
        population[i].domination_count = 0;
        population[i].dominated_solutions.clear();
        for (int j = 0; j < population.size(); ++j) {
            if (i == j) continue;

            bool p_dominates_q = (population[i].time < population[j].time && population[i].spill <= population[j].spill) ||
                (population[i].time <= population[j].time && population[i].spill < population[j].spill);
            bool q_dominates_p = (population[j].time < population[i].time && population[j].spill <= population[i].spill) ||
                (population[j].time <= population[i].time && population[j].spill < population[i].spill);

            if (p_dominates_q) {
                population[i].dominated_solutions.push_back(j);
            }
            else if (q_dominates_p) {
                population[i].domination_count++;
            }
        }
        if (population[i].domination_count == 0) {
            population[i].rank = 1;
            fronts[0].push_back(i);
        }
    }

    int front_idx = 0;
    while (front_idx < fronts.size() && !fronts[front_idx].empty()) {
        std::vector<int> next_front;
        for (int p_idx : fronts[front_idx]) {
            for (int q_idx : population[p_idx].dominated_solutions) {
                population[q_idx].domination_count--;
                if (population[q_idx].domination_count == 0) {
                    population[q_idx].rank = front_idx + 2;
                    next_front.push_back(q_idx);
                }
            }
        }
        front_idx++;
        if (!next_front.empty()) fronts.push_back(next_front);
    }

    for (auto& ind : population) ind.crowding_distance = 0.0;

    for (const auto& front : fronts) {
        if (front.size() < 3) {
            for (int idx : front) population[idx].crowding_distance = std::numeric_limits<double>::infinity();
            continue;
        }

        auto sorted_by_time = front;
        std::sort(sorted_by_time.begin(), sorted_by_time.end(), [&](int a, int b) {
            return population[a].time < population[b].time;
            });

        auto sorted_by_spill = front;
        std::sort(sorted_by_spill.begin(), sorted_by_spill.end(), [&](int a, int b) {
            return population[a].spill < population[b].spill;
            });

        population[sorted_by_time.front()].crowding_distance = std::numeric_limits<double>::infinity();
        population[sorted_by_time.back()].crowding_distance = std::numeric_limits<double>::infinity();
        population[sorted_by_spill.front()].crowding_distance = std::numeric_limits<double>::infinity();
        population[sorted_by_spill.back()].crowding_distance = std::numeric_limits<double>::infinity();

        double time_range = population[sorted_by_time.back()].time - population[sorted_by_time.front()].time;
        double spill_range = population[sorted_by_spill.back()].spill - population[sorted_by_spill.front()].spill;

        if (time_range > 0) {
            for (size_t i = 1; i < sorted_by_time.size() - 1; ++i) {
                population[sorted_by_time[i]].crowding_distance += (population[sorted_by_time[i + 1]].time - population[sorted_by_time[i - 1]].time) / time_range;
            }
        }
        if (spill_range > 0) {
            for (size_t i = 1; i < sorted_by_spill.size() - 1; ++i) {
                population[sorted_by_spill[i]].crowding_distance += (population[sorted_by_spill[i + 1]].spill - population[sorted_by_spill[i - 1]].spill) / spill_range;
            }
        }
    }
}